import numpy as np
import cv2
import random
import math
import json


"""
# Template:
# <class ObjDetDataAug(object)>
class ObjDetDataAug(object):
    # <method __init__>
    def __init__(self):
        pass
    # <method __init__>

    # <method __call__>
    def __call__(self, *args, **kwargs):
        image_np = args[0]["image"]
        boxes_np = args[0]["target"]
        raise NotImplementedError("")
        return {"image": out_img_np, "target": out_boxes_np}
    # <method __call__>
# <class ObjDetDataAug(object)>
"""


# <class ObjDetDataAugRandomHorizontalFlip(object)>
class ObjDetDataAugRandomHorizontalFlip(object):
    # <method __init__>
    def __init__(self):
        pass
    # <method __init__>

    # <method __call__>
    def __call__(self, *args, **kwargs):
        image_np = args[0]["image"]
        boxes_np = args[0]["target"]
        # random flip:
        if random.randint(0, 1) > 0:
            out_img_np = cv2.flip(image_np, 1)
            out_boxes_np = boxes_np.copy()
            out_boxes_np[..., 1] = image_np.shape[1] - 1 - boxes_np[..., 1] # cx_flipped = img_w - 1 - cx
        else:
            out_img_np = image_np
            out_boxes_np = boxes_np
        # end-if
        return {"image": out_img_np, "target": out_boxes_np}
    # <method __call__>
# <class ObjDetDataAugRandomHorizontalFlip(object)>


# <class ObjDetDataAugRandomVerticalFlip(object)>
class ObjDetDataAugRandomVerticalFlip(object):
    # <method __init__>
    def __init__(self):
        pass
    # <method __init__>

    # <method __call__>
    def __call__(self, *args, **kwargs):
        image_np = args[0]["image"]
        boxes_np = args[0]["target"]
        # random flip:
        if random.randint(0, 1) > 0:
            out_img_np = cv2.flip(image_np, 0)
            out_boxes_np = boxes_np.copy()
            out_boxes_np[..., 0] = image_np.shape[0] - 1 - boxes_np[..., 0] # cy_flipped = img_h - 1 - cy
        else:
            out_img_np = image_np
            out_boxes_np = boxes_np
        # end-if
        return {"image": out_img_np, "target": out_boxes_np}
    # <method __call__>
# <class ObjDetDataAugRandomVerticalFlip(object)>


# <class ObjDetDataAugRandomRotate(object)>
class ObjDetDataAugRandomRotate(object):
    # <method __init__>
    def __init__(self, max_rotate_angle = 45):
        self._max_rotate_angle = max_rotate_angle
        pass
    # <method __init__>

    # <method __call__>
    def __call__(self, *args, **kwargs):
        image_np = args[0]["image"]
        boxes_np = args[0]["target"]
        # # start
        angle = random.randint(0, int(self._max_rotate_angle * 2)) - int(self._max_rotate_angle)
        M = cv2.getRotationMatrix2D((image_np.shape[1]//2, image_np.shape[0]//2), angle, 1.0)
        vcos = np.abs(M[0, 0])
        vsin = np.abs(M[0, 1])
        rotated_img_h = int(image_np.shape[0] * vcos + image_np.shape[1] * vsin)
        rotated_img_w = int(image_np.shape[0] * vsin + image_np.shape[1] * vcos)
        M[0, 2] += (rotated_img_w // 2) - (image_np.shape[1] // 2)
        M[1, 2] += (rotated_img_h // 2) - (image_np.shape[0] // 2)
        # get rotated out_img_np
        out_img_np = cv2.warpAffine(src=image_np, M=M, dsize=(rotated_img_w, rotated_img_h), borderValue=(128, 128, 128))
        # get rotated out_boxes_np
        out_boxes_np = boxes_np.copy()
        out_boxes_np[:, 0] = M[1, 0] * boxes_np[:, 1] + M[1, 1] * boxes_np[:, 0] + M[1, 2]
        out_boxes_np[:, 1] = M[0, 0] * boxes_np[:, 1] + M[0, 1] * boxes_np[:, 0] + M[0, 2]    
        out_boxes_np[:, 2] = boxes_np[:, 2] * vcos + boxes_np[:, 3] * vsin
        out_boxes_np[:, 3] = boxes_np[:, 2] * vsin + boxes_np[:, 3] * vcos
        # get dict
        return {"image": out_img_np, "target": out_boxes_np}
    # <method __call__>
# <class ObjDetDataAugRandomRotate(object)>



if __name__ == "__main__":
    
    pass